import random as rd
import numpy as np


def Partition(a, p: int, r: int):
    j = r
    i = p + 1
    x = a[p]
    while True:
        while a[i] <= x and i < r:  # a[i]<=x   左边，第一次停在a[i]>x处
            i = i + 1
        while a[j] >= x and j > p:  # a[j]>=x   右边，第一次停在a[j]<x处
            j = j - 1
        if i >= j:  # 这时如果i>=j,说明数组已经分成左右两块
            break
        a[i], a[j] = a[j], a[i]
    a[p] = a[j]
    a[j] = x
    # print(a)
    return j


def myPartition(a, p, r, x):
    a[x], a[p] = a[p], a[x]
    return Partition(a, p, r)


# 随机选取基准，改进算法
def RandomizedPartiton(a, p, r):
    i = rd.randint(p, r)
    a[i], a[p] = a[p], a[i]
    return Partition(a, p, r)


def RandomizedQuickSort(a, p, r):
    if p < r:
        q = RandomizedPartiton(a, p, r)
        RandomizedQuickSort(a, p, q)
        RandomizedQuickSort(a, q + 1, r)


# 方法二: 最坏情况下也能O(n)
def Select(a, p: int, r: int, k: int):
    if r - p < 75:
        RandomizedQuickSort(a, p, r)
        return p + k - 1
        # return a[p + k - 1]
    else:
        for item in range((r - p - 4) // 5 + 1):  # 将其按5个一组分组从a[p+i*5]到 a[p+i*5+4]
            x = Select(a, p, p + (r - p - 4) // 5, (r - p - 4) // 10)  # 找出中位数的中位数
            i = myPartition(a, p, r, x)
            j = i - p + 1
            if k <= j:
                return Select(a, p, i, k)
            else:
                return Select(a, i + 1, r, k - j)


if __name__ == "__main__":
    # a = np.empty([100], dtype=int)
    a = np.random.random_integers(-90, 900, size=800)
    print("排序前", a, "长度:", len(a))

    k = -1
    # 当k小于1后，由于数组的特性会变成倒着来的，如第0小会是最大值
    if len(a) >= k > 0:
        ans = Select(a, 0, len(a) - 1, k)
        print("second way 第", k, "小：", a[ans])
    else:
        print("error")
    RandomizedQuickSort(a, 0, len(a) - 1)
    print("真实值", a[k-1])
